import os
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from dotenv import load_dotenv
from loguru import logger

# 加载环境变量
load_dotenv()

# 从环境变量获取密钥，如果不存在则生成一个新的
ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY","U8Jq3WBL6qJ2wvKk1ZAx0d5FfYgR7tEhYsTmNcWuOos=")
if not ENCRYPTION_KEY:
    # 为了演示，这里直接生成一个密钥
    # 生产环境应通过安全的方式管理密钥
    logger.warning("环境变量中未找到ENCRYPTION_KEY，将生成一个新的密钥")
    
    # 生成密钥
    salt = os.urandom(16)
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(b"data_source_password"))
    ENCRYPTION_KEY = key.decode()
    
    # 输出密钥以便管理员设置环境变量
    logger.info(f"生成的新密钥为：{ENCRYPTION_KEY}，请设置到环境变量ENCRYPTION_KEY中")
else:
    logger.info("已从环境变量加载ENCRYPTION_KEY")

# 创建Fernet实例
try:
    cipher = Fernet(ENCRYPTION_KEY.encode() if isinstance(ENCRYPTION_KEY, str) else ENCRYPTION_KEY)
except Exception as e:
    logger.error(f"初始化加密器失败：{e}")
    # 为了保证程序正常启动，使用临时密钥
    temp_key = Fernet.generate_key()
    cipher = Fernet(temp_key)
    logger.warning("使用临时密钥进行加密，所有现有加密数据将无法解密！请尽快设置正确的ENCRYPTION_KEY环境变量。")


def encrypt_text(text):
    """加密文本
    
    Args:
        text: 要加密的文本
        
    Returns:
        加密后的文本
    """
    if not text:
        return None
    
    try:
        if isinstance(text, str):
            text = text.encode()
        
        encrypted_text = cipher.encrypt(text)
        return encrypted_text.decode()
    except Exception as e:
        logger.error(f"加密文本失败: {e}")
        return None


def decrypt_text(encrypted_text):
    """解密文本
    
    Args:
        encrypted_text: 加密后的文本
        
    Returns:
        解密后的文本
    """
    if not encrypted_text:
        return None
    
    try:
        if isinstance(encrypted_text, str):
            encrypted_text = encrypted_text.encode()
            
        decrypted_text = cipher.decrypt(encrypted_text)
        return decrypted_text.decode()
    except Exception as e:
        logger.error(f"解密文本失败: {e}")
        return None 